package com.yupi.springbootinit.api;

import java.io.BufferedReader;
import java.io.DataOutputStream;
import java.io.InputStreamReader;
import java.net.HttpURLConnection;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import com.google.gson.Gson;
import com.google.gson.JsonObject;

public class OpenAIChatClient {
    private static final String API_URL = "https://api.qyfxw.cn/v1/chat/completions";
    private final String apiKey;

    public OpenAIChatClient(String apiKey) {
        this.apiKey = apiKey;
    }

    public String chat(String message) throws Exception {
        URL url = new URL(API_URL);
        HttpURLConnection connection = (HttpURLConnection) url.openConnection();
        connection.setRequestMethod("POST");
        connection.setRequestProperty("Content-Type", "application/json");
        connection.setRequestProperty("Authorization", "Bearer " + apiKey);
        connection.setDoOutput(true);

        Gson gson = new Gson();
        JsonObject requestBody = new JsonObject();
        requestBody.addProperty("model", "gpt-3.5-turbo");
        JsonObject messageObj = new JsonObject();
        messageObj.addProperty("role", "user");
        messageObj.addProperty("content", message);
        requestBody.add("messages", gson.toJsonTree(new JsonObject[]{messageObj}));

        byte[] input = requestBody.toString().getBytes(StandardCharsets.UTF_8);
        try (DataOutputStream dos = new DataOutputStream(connection.getOutputStream())) {
            dos.write(input, 0, input.length);
        }

        int responseCode = connection.getResponseCode();
        StringBuilder response = new StringBuilder();
        try (BufferedReader br = new BufferedReader(
                new InputStreamReader(responseCode == 200 ? connection.getInputStream() : connection.getErrorStream(), StandardCharsets.UTF_8))) {
            String responseLine;
            while ((responseLine = br.readLine()) != null) {
                response.append(responseLine.trim());
            }
        }

        if (responseCode != 200) {
            throw new Exception("HTTP request failed with code: " + responseCode + ", response: " + response.toString());
        }

        JsonObject responseJson = gson.fromJson(response.toString(), JsonObject.class);
        return responseJson.getAsJsonArray("choices").get(0).getAsJsonObject().getAsJsonObject("message").get("content").getAsString();
    }
}